package com.company.nuwa.common.handler;

import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * <p>线程池统一收口处理</p >
 *
 * @author youyoukele
 * @version 1.0
 * @date 2021/03/08 16:35
 */
@Slf4j
public class ThreadPoolHandler {

    /**
     * 构建线程池参数
     *
     * @param prefix           线程池名称前缀
     * @param keepAliveSeconds 保持空闲时间
     * @param queueCapacity    队列大小
     * @param corePoolSize     核心线程池大小
     * @param maxPoolSize      最大线程池大小
     * @param strategy         队列满之后的处理策略
     * @return org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor
     **/
    public static ThreadPoolTaskExecutor build(String prefix, int keepAliveSeconds, int queueCapacity
            , int corePoolSize, int maxPoolSize, RejectedHandlerStrategy strategy) {
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        threadPoolTaskExecutor.setThreadNamePrefix(prefix);
        threadPoolTaskExecutor.setCorePoolSize(corePoolSize);
        threadPoolTaskExecutor.setMaxPoolSize(maxPoolSize);
        threadPoolTaskExecutor.setKeepAliveSeconds(keepAliveSeconds);
        threadPoolTaskExecutor.setQueueCapacity(queueCapacity);
        threadPoolTaskExecutor.setRejectedExecutionHandler(strategy.getHandler());
        threadPoolTaskExecutor.setTaskDecorator(new ThreadTaskDecorator());
        threadPoolTaskExecutor.setWaitForTasksToCompleteOnShutdown(true);
        // 线程池关闭时，超过60秒，如果还有线程未销毁，则强制销毁，并非一直阻塞
        threadPoolTaskExecutor.setAwaitTerminationSeconds(60);
        return threadPoolTaskExecutor;
    }

    /**
     * 拒绝策略
     */
    @Getter
    @AllArgsConstructor
    public enum RejectedHandlerStrategy {
        /**
         * caller
         */
        CALLER("caller", new CallerRunsPolicy()),
        /**
         * abort
         */
        ABORT("abort", new AbortPolicy()),
        /**
         * discard
         */
        DISCARD("discard", new DiscardPolicy()),
        ;

        /**
         * 编码
         */
        private final String code;
        /**
         * 处理方式
         */
        private final RejectedExecutionHandler handler;
    }

    /**
     * A handler for rejected tasks that runs the rejected task
     * directly in the calling thread of the {@code execute} method,
     * unless the executor has been shut down, in which case the task
     * is discarded.
     */
    public static class CallerRunsPolicy implements RejectedExecutionHandler {
        /**
         * Creates a {@code CallerRunsPolicy}.
         */
        public CallerRunsPolicy() {
        }

        /**
         * Executes task r in the caller's thread, unless the executor
         * has been shut down, in which case the task is discarded.
         *
         * @param r the runnable task requested to be executed
         * @param e the executor attempting to execute this task
         */
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            if (!e.isShutdown()) {
                log.warn("CallerRunsPolicy Task " + r.toString() +
                        " caller run from " +
                        e.toString());
                r.run();
            }
        }
    }

    /**
     * A handler for rejected tasks that throws a
     * {@code RejectedExecutionException}.
     */
    public static class AbortPolicy implements RejectedExecutionHandler {
        /**
         * Creates an {@code AbortPolicy}.
         */
        public AbortPolicy() {
        }

        /**
         * Always throws RejectedExecutionException.
         *
         * @param r the runnable task requested to be executed
         * @param e the executor attempting to execute this task
         * @throws RejectedExecutionException always
         */
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            throw new RejectedExecutionException("Task " + r.toString() +
                    " rejected from " +
                    e.toString());
        }
    }

    /**
     * A handler for rejected tasks that silently discards the
     * rejected task.
     */
    public static class DiscardPolicy implements RejectedExecutionHandler {
        /**
         * Creates a {@code DiscardPolicy}.
         */
        public DiscardPolicy() {
        }

        /**
         * Does nothing, which has the effect of discarding task r.
         *
         * @param r the runnable task requested to be executed
         * @param e the executor attempting to execute this task
         */
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor e) {
            log.error("DiscardPolicy Task " + r.toString() +
                    " discard from " +
                    e.toString());
        }
    }

}
